from torch.utils.data import Dataset
from torch.utils.data.dataset import T_co
from transformers import BertTokenizer
from utils.plus_function import get_bin_str_list_range_n
from utils.raw_data1 import RawData
from exp02.instance import SentenceInstance, TrainingInstance, TestInstance


class TrainDataSet(Dataset):  # train dataset 和 test dataset 应该是不一样的
    def __init__(self, raw_data: RawData, tokenizer: BertTokenizer):
        plus_text = "Find the first aspect term and corresponding opinion term in the text"
        plus_meta_sentence_instance = SentenceInstance(plus_text, [], tokenizer)
        self.training_instances = []

        def generate_train_instance_list(meta_ins: SentenceInstance):
            n_triplet = len(meta_ins.triplets)
            bin_str_list = get_bin_str_list_range_n(n_triplet)
            res = []
            for bin_str in bin_str_list:
                aspect_i = bin_str.find("1")
                if aspect_i < 0:
                    continue
                    pass
                n_opinion = len(meta_ins.triplets[aspect_i]["opinion_span_list"])
                opinion_bin_str_list = get_bin_str_list_range_n(n_opinion)
                res.extend(
                    [TrainingInstance(
                        meta_ins,
                        plus_meta_sentence_instance,
                        bin_str,
                        opinion_bin_str
                    ) for opinion_bin_str in opinion_bin_str_list]
                )
            return res

        for meta_instance in raw_data:
            current_train_instance_list = generate_train_instance_list(meta_instance)
            self.training_instances.extend(current_train_instance_list)
            pass
        pass

    def __getitem__(self, index) -> T_co:
        return self.training_instances[index]

    def __iter__(self):
        for elem in self.training_instances:
            yield elem
            pass
        pass

    def __len__(self) -> int:
        return len(self.training_instances)

    pass
